package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")

type GuessArgs struct {
	Name string `json:"name" jsonschema:"姓名"`
}

type GuessResp struct {
	Ans string `json:"ans" jsonschema:"答案"`
}

func GuessLike(ctx context.Context, session *mcp.ServerSession, params *mcp.CallToolParamsFor[GuessArgs]) (*mcp.CallToolResultFor[GuessResp], error) {
	return &mcp.CallToolResultFor[GuessResp]{
		Content: []mcp.Content{
			&mcp.TextContent{Text: ""},
		},
		StructuredContent: GuessResp{
			Ans: func() string {
				name := params.Arguments.Name
				ansList := []string{"香蕉", "苹果", "水蜜桃"}
				return fmt.Sprintf("%s喜欢%s", name, ansList[rand.Intn(3)])
			}(),
		},
	}, nil
}

type SumArgs struct {
	A int `json:"a" jsonschema:"第一个数"`
	B int `json:"b" jsonschema:"第二个数"`
}

type SumRes struct {
	Sum int `json:"sum" jsonschema:"两数之和"`
}

func Sum(ctx context.Context, session *mcp.ServerSession, params *mcp.CallToolParamsFor[SumArgs]) (*mcp.CallToolResultFor[SumRes], error) {
	res := SumRes{
		Sum: params.Arguments.A + params.Arguments.B,
	}
	resByte, _ := json.Marshal(res)
	return &mcp.CallToolResultFor[SumRes]{
		Content: []mcp.Content{
			&mcp.TextContent{Text: string(resByte)},
		},
		StructuredContent: res,
	}, nil
}

func PromptHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
	return &mcp.GetPromptResult{
		Description: "Code review prompt",
		Messages: []*mcp.PromptMessage{
			{Role: "user", Content: &mcp.TextContent{Text: "Say hi to " + params.Arguments["name"]}},
		},
	}, nil
}

func main() {
	flag.Parse()

	server := mcp.NewServer(&mcp.Implementation{Name: "MCP测试服务器", Title: "my mcp server", Version: "1.0.0"}, nil)
	mcp.AddTool(server, &mcp.Tool{Name: "guessLike", Description: "猜一个人喜欢什么"}, GuessLike)
	mcp.AddTool(server, &mcp.Tool{Name: "sum", Title: "两数求和", Description: "得到两个数的和"}, Sum)
	// server.AddPrompt(&mcp.Prompt{Name: "greet"}, PromptHi)
	// server.AddResource(&mcp.Resource{
	// 	Name:     "info",
	// 	MIMEType: "text/plain",
	// 	URI:      "embedded:info",
	// }, handleEmbeddedResource)

	if *httpAddr != "" {
		handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
			return server
		}, nil)
		log.Printf("MCP handler listening at %s", *httpAddr)
		http.ListenAndServe(*httpAddr, handler)
	} else {
		if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil {
			log.Printf("Server failed: %v", err)
		}
	}
}

var embeddedResources = map[string]string{
	"info": "This is the hello example server.",
}

func handleEmbeddedResource(ctx context.Context, session *mcp.ServerSession, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) {
	u, err := url.Parse(params.URI)
	if err != nil {
		return nil, err
	}
	if u.Scheme != "embedded" {
		return nil, fmt.Errorf("wrong scheme: %q", u.Scheme)
	}
	key := u.Opaque
	text, ok := embeddedResources[key]
	if !ok {
		return nil, fmt.Errorf("no embedded resource named %q", key)
	}
	return &mcp.ReadResourceResult{
		Contents: []*mcp.ResourceContents{
			{URI: params.URI, MIMEType: "text/plain", Text: text},
		},
	}, nil
}
